import source
from ai.audio_dataset.classify import data_handler
from ai.audo_model.classify import xvector
from ai.config.config import GxlNode
from ai.utils import utils_file, utils_model
import torch
import torch.nn as nn
from project.wav_classify.my_runner import GxlRunnerGxl


def main():
    config = GxlNode.get_config_from_yaml('./config.yaml')
    utils_model.set_random_seed(config.random.seed)
    data_dev_iter = data_handler.get_iter_by_jsonl('./output/accent_dev_little_100.jsonl', config)
    data_train_iter = data_handler.get_iter_by_json('./output/accent_train.json', config)
    classify_model = xvector.GxlClassifier()
    optim = torch.optim.Adam(classify_model.parameters(), lr=config.optim.lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optim, T_0=5, T_mult=2)
    loss_f = nn.CrossEntropyLoss()
    runner_man = GxlRunnerGxl(classify_model, optim, loss_f,
                              data_train_iter, config, data_dev_iter, scheduler, device=torch.device('cuda:0'))
    runner_man.run()


def inference():
    config = GxlNode.get_config_from_yaml('./config_inference.yaml')
    data_dev_little_iter = data_handler.get_iter_by_json('./output/accent_same.json', config)
    classify_model = xvector.GxlClassifier()
    optim = torch.optim.Adam(classify_model.parameters(), lr=config.optim.lr)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optim, T_0=5, T_mult=2)
    loss_f = nn.CrossEntropyLoss()
    runner_man = GxlRunnerGxl(classify_model, optim, loss_f,
                              None, config, data_dev_little_iter, scheduler, device=torch.device('cuda:0'))
    runner_man.inference(data_dev_little_iter, './output/inference1.jsonl')
    # a = runner_man.calculate_valid_loss()
    # print(a)


if __name__ == '__main__':
    inference()
    # main()